import torch


def sample_gaussian(m, v):
    sample = torch.randn(m.shape).cuda()
    z = m + (v ** 0.5) * sample
    return z


def kl_normal(qm, qv, pm, pv, yh):
    element_wise = 0.5 * (torch.log(pv) - torch.log(qv) + qv / pv + (qm - pm - yh).pow(2) / pv - 1)
    kl = element_wise.sum(-1)
    # print("log var1", qv)
    return kl


def get_exp_name(args):
    cf_option_str = [args.dataset]

    if args.lamda:
        additional_str = 'lamda%s' % args.lamda
        cf_option_str.append(additional_str)

    if args.lr and args.lr != 0.001:
        additional_str = 'lr%s' % args.lr
        cf_option_str.append(additional_str)

    if args.beta_z != 1:
        additional_str = 'betaz%s' % args.beta_z
        cf_option_str.append(additional_str)

    if args.encode_z:
        additional_str = 'encodez%s' % args.encode_z
        cf_option_str.append(additional_str)

    if args.contrastive_loss:
        additional_str = 'contra'
        cf_option_str.append(additional_str)

    if args.contrastive_loss and args.temperature:
        additional_str = 'T%s' % args.temperature
        cf_option_str.append(additional_str)

    # if args.weight_decay != 0.00:
    #     additional_str = 'weight_decay%s' % args.weight_decay
    #     cf_option_str.append(additional_str)

    return "-".join(cf_option_str)
